package main

import (
	"crypto/md5"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type RSACipher struct {
	pub *rsa.PublicKey
	pri *rsa.PrivateKey
}

func (this *RSACipher) Load(key string) error {
	f, err := os.Open(key)
	if err != nil {
		return err
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		return err
	}
	block, _ := pem.Decode(b)
	if block == nil {
		return errors.New("invalid rsa pem key")
	}

	switch {
	case strings.Contains(block.Type, "PRIVATE"):
		this.pri, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return err
		}
	case strings.Contains(block.Type, "PUBLIC"):
		pub, err := x509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return err
		}
		var ok bool
		this.pub, ok = pub.(*rsa.PublicKey)
		if !ok {
			return errors.New("invalid rsa public key")
		}
	default:
		return errors.New("unknow rsa pem key type")
	}

	return nil
}

func (this *RSACipher) Encryp(plaintext []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, this.pub, plaintext)
}

func (this *RSACipher) Decryp(ciphertext []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, this.pri, ciphertext)
}

func (this *RSACipher) Sign(hashed []byte) ([]byte, error) {
	return rsa.SignPKCS1v15(rand.Reader, this.pri, 0, hashed)
}

func (this *RSACipher) Verify(hashed []byte, sign []byte) error {
	return rsa.VerifyPKCS1v15(this.pub, 0, hashed, sign)
}

func main() {
	c := new(RSACipher)
	c.Load("public.pem")
	c.Load("private.pem")

	{
		b, err := c.Encryp([]byte("Hello"))
		if err != nil {
			panic(err)
		}
		fmt.Println(b)

		s, err := c.Decryp(b)
		if err != nil {
			panic(err)
		}
		fmt.Println(string(s))
	}

	{
		s := md5.Sum([]byte("Hello"))
		hashed := s[:]
		fmt.Println("hashed:", hashed)

		sign, err := c.Sign(hashed)
		if err != nil {
			panic(err)
		}

		err = c.Verify(hashed, sign)
		if err != nil {
			panic(err)
		}
		fmt.Println("verify success")

		err = c.Verify(hashed, sign[:len(sign)-3])
		if err != nil {
			fmt.Println(err)
			fmt.Println("use invalid sign success")
		}
	}
}
